import os
import json
import flask
import requests
from flask import Flask, request, jsonify, Response
from flask_cors import CORS  # 导入 CORS
from ollama import Client
import logging
import sys
from search import WebSearcher

# 添加日志配置
logging.basicConfig(
    level=logging.INFO,  # 设置日志级别为 INFO
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',  # 设置日志格式
    handlers=[
        logging.StreamHandler(sys.stdout),  # 输出到控制台
        logging.FileHandler('app.log', encoding='utf-8')  # 同时保存到文件
    ]
)

# 配置信息
SERPER_API_KEY = os.environ.get('SERPER_API_KEY', 'AIzaSyAr5OkcqVi1AfaF9_veGsGLoPtpPSXwgn8')
OLLAMA_HOST = 'http://127.0.0.1:11434'
PROXY_PORT = 8000

class OllamaSearchProxy:
    def __init__(self, serper_key, ollama_host):
        self.serper_key = serper_key
        self.ollama_client = Client(host=ollama_host)  # 在这里加载模型
        self.app = Flask(__name__)
        CORS(self.app)  # 启用 CORS
        self.logger = logging.getLogger(__name__)  # 使用上面配置的logger
        self.setup_routes()
        self.webSearcher = WebSearcher()

    def enhance_ollama_search_prompt(slef, query, enhanced_results):
        """
        3.根据获取的网页内容生成更精确的提示词
        :param query: 原始查询
        :param enhanced_results: 包含网页内容的搜索结果
        :return: 增强型提示词
        """
        prompt = f"针对问题： '{query}', 以下是相关网页内容：\n\n"
        
        for idx, result in enumerate(enhanced_results, 1):
            prompt += f"来源 {idx}: {result['title']}\n"
            prompt += f"链接: {result['link']}\n"
            prompt += f"摘要: {result['snippet']}\n"
            prompt += f"主要内容: {result['content']}...\n\n"
        
        prompt += "请基于这些详细信息，提供一个全面、准确的回答。注意整合不同来源的信息，并保持客观中立。"
        
        return prompt        

    def setup_routes(self):
        """
        3. 设置Flask路由
        """
        @self.app.route('/search', methods=['POST'])
        def web_search():
            data = request.get_json()
            query = data.get('query', '')            
            # 执行网络搜索
            search_result = self.webSearcher.search_web(query)
            return jsonify(search_result)

        @self.app.route('/chat', methods=['POST'])
        def chat_with_search():
            data = request.get_json()
            query = data.get('query', '')
            # model = data.get('model', 'qwen2.5:14b-instruct-q3_K_S')  # 默认使用通义千问
            model = 'qwen2.5:32b';
            self.logger.info(f"使用模型: {model}, 收到聊天请求，查询内容: {query}")
            
            # 执行网络搜索
            search_result = self.webSearcher.search_web(query)            
            if search_result['status'] == 'error':
                return jsonify({
                    'status': 'error',
                    'message': search_result['message']
                })

            # 爬取链接内容
            content_results = search.crawler_content(search_result);            
            # 生成增强型提示词
            enhanced_prompt = self.enhance_ollama_search_prompt(query, content_results)
            
            def generate_response():
                try:
                    # 调用Ollama API
                    response = requests.post('http://10.169.3.15:8088/chat-api/api/chat',
                        # 'http://localhost:11434/api/chat',
                        json={
                            'model': model,
                            'messages': [
                                {
                                    'role': 'system', 
                                    'content': '你是一个基于网络搜索结果的智能助手。请仔细分析提供的搜索信息，并给出准确的回答。'
                                },
                                {
                                    'role': 'user', 
                                    'content': enhanced_prompt
                                }
                            ]
                        },
                        stream=True  # 启用流式输出
                    )
                    
                    # 逐步返回内容
                    for line in response.iter_lines():
                        if line:
                            yield f"{line.decode('utf-8')}\n\n"  # 使用 Server-Sent Events 格式
                
                except Exception as e:
                    yield f"data: 调用Ollama API时发生错误：{str(e)}\n\n"

            return Response(generate_response(), mimetype='text/event-stream')

    def run(self):
        """
        4. 启动代理服务器
        """
        self.app.run(host='127.0.0.1', port=PROXY_PORT)

def main():
    # 创建并运行代理服务器
    proxy = OllamaSearchProxy(
        serper_key=SERPER_API_KEY, 
        ollama_host=OLLAMA_HOST
    )
    proxy.run()

if __name__ == '__main__':
    main()
